"""
Utility functions for data processing and manipulation
"""
import numpy as np
import torch
from scipy.signal import butter, lfilter
import matplotlib.pyplot as plt
import pandas as pd
import os


def butter_bandpass(lowcut, highcut, fs, order=5):
    """
    Design a Butterworth bandpass filter

    Args:
        lowcut (float): Low cut frequency in Hz
        highcut (float): High cut frequency in Hz
        fs (float): Sampling rate in Hz
        order (int): Filter order

    Returns:
        tuple: Filter coefficients (b, a)
    """
    nyq = 0.5 * fs
    low = lowcut / nyq
    high = highcut / nyq
    b, a = butter(order, [low, high], btype='band')
    return b, a


def butter_bandpass_filter(data, lowcut, highcut, fs, order=5):
    """
    Apply a Butterworth bandpass filter to the signal

    Args:
        data (numpy.ndarray): Input signal
        lowcut (float): Low cut frequency in Hz
        highcut (float): High cut frequency in Hz
        fs (float): Sampling rate in Hz
        order (int): Filter order

    Returns:
        numpy.ndarray: Filtered signal
    """
    b, a = butter_bandpass(lowcut, highcut, fs, order=order)
    y = lfilter(b, a, data)
    return y


def extract_frequency_bands(signal, fs):
    """
    Extract different frequency bands from the EEG signal

    Args:
        signal (numpy.ndarray): Input EEG signal
        fs (float): Sampling rate in Hz

    Returns:
        dict: Dictionary containing filtered signals for each frequency band
    """
    # Define frequency bands
    bands = {
        'delta': (0.5, 4),  # Delta band: 0.5-4 Hz
        'theta': (4, 8),  # Theta band: 4-8 Hz
        'alpha': (8, 13),  # Alpha band: 8-13 Hz
        'sigma': (12, 16),  # Sigma band: 12-16 Hz (sleep spindles)
        'beta': (16, 30),  # Beta band: 16-30 Hz
    }

    # Filter signal for each band
    filtered_bands = {}
    for band_name, (low, high) in bands.items():
        filtered_bands[band_name] = butter_bandpass_filter(signal, low, high, fs)

    return filtered_bands


def normalize_signal(signal, method='minmax'):
    """
    Normalize the EEG signal

    Args:
        signal (numpy.ndarray): Input EEG signal
        method (str): Normalization method ('minmax', 'zscore')

    Returns:
        numpy.ndarray: Normalized signal
    """
    if method == 'minmax':
        # Min-max normalization to [0, 1]
        min_val = np.min(signal)
        max_val = np.max(signal)

        # Avoid division by zero
        if max_val == min_val:
            return np.zeros_like(signal)

        return (signal - min_val) / (max_val - min_val)

    elif method == 'zscore':
        # Z-score normalization
        mean = np.mean(signal)
        std = np.std(signal)

        # Avoid division by zero
        if std == 0:
            return np.zeros_like(signal)

        return (signal - mean) / std

    else:
        raise ValueError(f"Unknown normalization method: {method}")


def create_batches_for_sequence(data, seq_length, step=1):
    """
    Create batches of sequences from the data

    Args:
        data (numpy.ndarray): Input data
        seq_length (int): Length of each sequence
        step (int): Step size between sequences

    Returns:
        numpy.ndarray: Batched sequences
    """
    n_samples = len(data)
    n_batches = (n_samples - seq_length) // step + 1

    batches = []
    for i in range(0, n_samples - seq_length + 1, step):
        batch = data[i:i + seq_length]
        batches.append(batch)

    return np.array(batches)


def one_hot_encode(y, num_classes=5):
    """
    One-hot encode the labels

    Args:
        y (numpy.ndarray): Input labels
        num_classes (int): Number of classes

    Returns:
        numpy.ndarray: One-hot encoded labels
    """
    return np.eye(num_classes)[y]


def save_results_to_csv(results, filename, mode='w'):
    """
    Save results to CSV file

    Args:
        results (dict): Dictionary containing results
        filename (str): Output filename
        mode (str): File open mode ('w', 'a')
    """
    if not os.path.exists(os.path.dirname(filename)):
        os.makedirs(os.path.dirname(filename))

    df = pd.DataFrame([results])
    df.to_csv(filename, mode=mode, header=(mode == 'w'), index=False)


def plot_confusion_matrix(cm, class_names, save_path=None):
    """
    Plot confusion matrix

    Args:
        cm (numpy.ndarray): Confusion matrix
        class_names (list): List of class names
        save_path (str, optional): Path to save the figure
    """
    plt.figure(figsize=(10, 8))

    # Normalize the confusion matrix
    cm_norm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]

    plt.imshow(cm_norm, interpolation='nearest', cmap=plt.cm.Blues)
    plt.title('Confusion Matrix')
    plt.colorbar()

    tick_marks = np.arange(len(class_names))
    plt.xticks(tick_marks, class_names, rotation=45)
    plt.yticks(tick_marks, class_names)

    # Add normalized values inside the plot
    thresh = cm_norm.max() / 2.
    for i in range(cm_norm.shape[0]):
        for j in range(cm_norm.shape[1]):
            plt.text(j, i, f'{cm[i, j]}\n({cm_norm[i, j]:.2f})',
                     horizontalalignment="center",
                     color="white" if cm_norm[i, j] > thresh else "black")

    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label')

    if save_path:
        plt.savefig(save_path)
        plt.close()
    else:
        plt.show()


def plot_training_history(history, save_path=None):
    """
    Plot training history

    Args:
        history (dict): Dictionary containing training history
        save_path (str, optional): Path to save the figure
    """
    plt.figure(figsize=(15, 5))

    # Plot training & validation accuracy
    plt.subplot(1, 2, 1)
    plt.plot(history['train_acc'], label='Train')
    plt.plot(history['val_acc'], label='Validation')
    plt.title('Model accuracy')
    plt.ylabel('Accuracy')
    plt.xlabel('Epoch')
    plt.legend()

    # Plot training & validation loss
    plt.subplot(1, 2, 2)
    plt.plot(history['train_loss'], label='Train')
    plt.plot(history['val_loss'], label='Validation')
    plt.title('Model loss')
    plt.ylabel('Loss')
    plt.xlabel('Epoch')
    plt.legend()

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path)
        plt.close()
    else:
        plt.show()


def calculate_class_weights(labels, num_classes=5):
    """
    Calculate class weights for handling imbalance

    Args:
        labels (numpy.ndarray): Input labels
        num_classes (int): Number of classes

    Returns:
        numpy.ndarray: Class weights
    """
    class_counts = np.bincount(labels, minlength=num_classes)
    total_samples = np.sum(class_counts)

    # Calculate weights inversely proportional to class frequencies
    weights = total_samples / (class_counts * num_classes)

    return weights